package decoy_free_fdr_estimation;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Random;

/**
 *
 * @author Ben Teng
 */
public class Decoy_free_FDR_estimation {

    private static int N;   //number of protein
    private static int M;   //number of peptide
    private static double Y[];//original score vector
    private  static ArrayList<String> proteins = new ArrayList<>(); //dictionary of protein
    private  static ArrayList<String> peptides = new ArrayList<>(); //dictionary of peptide
    private  static ArrayList<Double> scores = new ArrayList<>();   //dictionary of score
    private static double Pvalue[];
    private static double fdr[];
    private int peptide_degree_count[];  //degree of each peptide
    private int protein_degree_count []; //degree of each protein
    private  int edgecount;   //edge number  
    private  ArrayList<ArrayList<Integer>> graph_protein=new  ArrayList<>();
    
     public void initial(String fx, String fy) 
    {
        load_file(fx); 
        loadY(fy);   
    }
  
     //load bipartite graph
    public void load_file(String filename) 
    {   
        ArrayList<String> protein = new ArrayList<>();  
        ArrayList<String> peptide = new ArrayList<>(); 
        ArrayList<Double> score = new ArrayList<>();
        String s ;
        try {
            BufferedReader input = new BufferedReader(new FileReader(filename));
            while ((s = input.readLine()) != null) {        
                String info[] = s.split("\t");
                protein.add(info[1]);
                peptide.add(info[0]);
                score.add(Double.parseDouble(info[2]));
            }
            input.close();
        } catch (IOException | NumberFormatException e) {
            System.out.println("ERROR: can not load data");
        }

        
        for (int i = 0; i < protein.size(); i++) 
        {
            if (proteins.contains(protein.get(i))) 
                continue; 
            else 
                proteins.add(protein.get(i));
        }
        
        for (int i = 0; i < peptide.size(); i++) 
       {
            if (peptides.contains(peptide.get(i))) 
            {
                if(scores.get(peptides.indexOf(peptide.get(i)))<score.get(i))
                     scores.set(peptides.indexOf(peptide.get(i)),score.get(i));
            } 
            else 
            {
                peptides.add(peptide.get(i));
                scores.add(score.get(i));
            }
        }

        N = proteins.size();
        M = peptides.size();
        Y = new double[N];
        double [][]X = new double[proteins.size()][peptides.size()];
        for (int i = 0; i < protein.size(); i++) 
        {            
            X[proteins.indexOf(protein.get(i))][peptides.indexOf(peptide.get(i))] = 1;        
        }
        
        
        peptide_degree_count=new int[M];
        protein_degree_count =new int[N];
        
        for(int i=0;i<proteins.size();i++)
        {
            for(int j=0;j<peptides.size();j++)
            {
                if(X[i][j]==1)
                {
                    protein_degree_count[i]++;
                    peptide_degree_count[j]++;
                    edgecount++;
                }
                
            }
        }
        

        for(int i=0;i<proteins.size();i++)
         {
             ArrayList<Integer> temp=new ArrayList<>();   //对与某一 protein 存在edge 的 peptide
            for(int j=0;j<peptides.size();j++)
            {
                  if(X[i][j]!=0)
                  {
                      temp.add(j);
                  }
             }
             graph_protein.add(temp);
         }
        
    }
  
  //get original scores
    public void  loadY(String fy)
    {
         String s;
         ArrayList<String> FY = new ArrayList<>();         //protein 
         ArrayList<String> FYS = new ArrayList<>();         //score
         try {
            BufferedReader input1 = new BufferedReader(new FileReader(fy));
            while ((s = input1.readLine()) != null) {
                String info[] = s.split("\t");
                FYS.add(info[1]);
                FY.add(info[0]);
            }
            input1.close();
        } catch (Exception e) {
            System.out.println("ERROR: can not load score");
            System.exit(0);
        }
       
        for (int i = 0; i < N; i++)
        {
            for (int j = 0; j < N; j++)
            {
                if(proteins.get(i).contains(FY.get(j)))
                {
                    Y[i] = Double.parseDouble(FYS.get(j));
                    break;
                }
            }
        }
    }
 

   public  ArrayList<ArrayList<Integer>> graph_generation()
   {
      ArrayList<ArrayList<Integer>> random_graph_protein=new  ArrayList<>();
      for(int i=0;i<N;i++)
      {
          ArrayList<Integer> temp=new ArrayList<>();   //对与某一 protein 存在edge 的 peptide
          random_graph_protein.add(temp);
      } 
      Random r = new Random(System.currentTimeMillis());
      
      int protein_degree_count_temp[]=new int[N];
      int peptide_degree_count_temp[]=new int[M];
      ArrayList<Edge> edgelist=new ArrayList<>();
      
      for(int i=0;i<N;i++)
          protein_degree_count_temp[i]=protein_degree_count[i];
      for(int i=0;i<M;i++)
          peptide_degree_count_temp[i]=peptide_degree_count[i];
      
      int y[][]=new int[N][M];
      for(int i=0;i<edgecount;i++)
      { 
           int t=r.nextInt(N);              //the selected protein
           while(protein_degree_count_temp[t]<=0)
               t=r.nextInt(N);
    
           int flag=0;
           int p;                          //the selected peptide
           for(int pp=0;pp<M;pp++)
           {
               if(peptide_degree_count_temp[pp]>0&&y[t][pp]!=1)
               {
                   flag=0;
                   break;
               }
               flag=1;    
           }
           
           if(flag==0)
           {
               p=r.nextInt(M);
               while(peptide_degree_count_temp[p]<=0||y[t][p]==1)
                   p=r.nextInt(M);            
           }
           else
           {
                   int et=r.nextInt(edgelist.size());
                    int first=edgelist.get(et).first;
                    int second=edgelist.get(et).second;
                    while(y[t][second]==1)
                    {
                         et=r.nextInt(edgelist.size());
                         first=edgelist.get(et).first;
                         second=edgelist.get(et).second;
                    }
                    protein_degree_count_temp[first]++;
                    peptide_degree_count_temp[second]++; 
                     y[first][second]=0;
                    p=second;
                    edgelist.remove(et); 
                    i--;    
               
           }
           y[t][p]=1;
           Edge edge=new Edge();
           edge.first=t;
           edge.second=p;
           edgelist.add(edge);
           peptide_degree_count_temp[p]--;
           protein_degree_count_temp[t]--;
          
      } 
      
      for(int i=0;i<edgelist.size();i++)
      {
          int first=edgelist.get(i).first;
          int  second=edgelist.get(i).second;
          random_graph_protein.get(first).add(second);
      }
      return random_graph_protein;
  }
 
   
  //get coefficients vector beta
    public double[] lasso(double Y[]) 
    {
        
        double L=0;
        double Beta[] = new double[M]; 
        double A[] = new double[M];
        double B1[][] = new double[M][M];
        double C[] = new double[M];

       for (int j = 0; j < graph_protein.size(); j++) 
       {   
                for(int i=0;i<graph_protein.get(j).size();i++)
                {
                    int p=graph_protein.get(j).get(i);
                    A[p] += scores.get(p) * Y[j];
                    for(int k=0;k<graph_protein.get(j).size();k++)
                    {
                        int p1=graph_protein.get(j).get(k);
                        B1[p][p1]+= scores.get(p)*scores.get(p1);
                    }
                    C[p]+=scores.get(p)*scores.get(p);
                }       
        }
      
         ArrayList<Integer> AL = new ArrayList<>();
         int t=0;    
          while (AL.size()<M&&t<1000)
          {
              t++;     
              for (int i = 0; i < M; i++) 
             {
                 if(AL.contains(i))
                     continue;
                double B = 0;
                for (int j = 0; j < M; j++)
                        B += Beta[j] * B1[i][j];
   
               double temp  = Beta[i];  
               if(L>=Math.abs(A[i]-B))
                   Beta[i]=0;
                else if((A[i]-B)<0)
                     Beta[i] = (A[i] - B + L) / C[i] + Beta[i]; 
                else
                   Beta[i] = (A[i] - B - L) / C[i]+ Beta[i]; 
               if(Beta[i]>1)
                   Beta[i]=1;
               if(Beta[i]<0)
                   Beta[i]=0;
               
                if (Math.abs( Beta[i] - temp)>Math.pow(10, -5))
                {
                      AL.add(i);
                }
                else
                    Beta[i]=temp;
            }   
          }        
        return Beta;
    }
  
    
    public double[] linear(double[] Beta, ArrayList<ArrayList<Integer>> graph)
    {
           
        double Y1[] = new double[N];  
        
        for (int i = 0; i < N; i++)
        {
            Y1[i]=0;
            for (int j = 0; j < graph.get(i).size(); j++) 
            {
                int k=graph.get(i).get(j);
                Y1[i] += Beta[k] * scores.get(k);
            }
          if(Y1[i]>1)
              Y1[i]=1;
           if(Y1[i]<0)
               Y1[i]=0;
        }
        
        return Y1;
    }  
     
      
    public double[]compute_pvalue(int times,double Y[]) 
    {
        double Pvalue[] = new double[N];
        double Beta[]; 
        double[] tempY=new double[N];
        ArrayList<ArrayList<Integer>> random_graph_protein;
        Beta = lasso(Y);    
        for (int ii = 0; ii <times; ii++) 
        {
           random_graph_protein=graph_generation();
           tempY = linear(Beta,random_graph_protein);  
           for (int j = 0; j < N; j++)
          {     
              for(int k=0;k<N;k++)
              {
                  if(tempY[j]>Y[k])          
                  {
                       Pvalue[k]++;
                  }
              }
           }
        }  
       for (int i = 0; i < N; i++) 
       {    
           Pvalue[i] = Pvalue[i] /(N*times); 
       }    
        return Pvalue;
    }

    
    public double[] estimate_FDR(double PI_0, double Pvalue[]) 
    {
        double fdr[] = new double[N];
        int R[] = new int[N];
        for (int i = 0; i < N; i++) 
        {
            for (int j = 0; j < N; j++) 
            {
                if (Pvalue[j] <= Pvalue[i])
                {
                    R[i]++;
                }
            }
        }
        for (int i = 0; i < N; i++) 
        {
            if (Pvalue[i] == 0) 
            {
                fdr[i] = 0;
            } 
            else 
            {
                fdr[i] = PI_0 * Pvalue[i] * N / (R[i] * (1 - (Math.pow(1 - Pvalue[i], N))));
            }
        }
        return fdr;
    }

       
    public double choose_pi0(double Pvalue[])
    {
        double pi0 = 1;
        double p[]=new double[20];
        double pp[]=new double[20];
        double lambda[]=new double[20]; 
        double max=0;
        double minn=1;
        for(int i=0;i<N;i++)
        {
            if(Pvalue[i]>max)
                max=Pvalue[i];
            if(Pvalue[i]<minn)
                minn=Pvalue[i];
        }
        
        double mark=(max-minn)/21;
        for(int i=0;i<20;i++)
            lambda[i]=mark*i;
        
        
        for(int i=0;i<20;i++)
        {
            int c=0;
            for(int j=0;j<N;j++)
                if(Pvalue[j]>lambda[i])
                    c++;      
            pp[i]=c/((1-lambda[i])*N);         
            if(pi0>pp[i])
                pi0=pp[i];      
        }  
        
        double tempp[]=new double [N];
        for(int i=0;i<10000;i++)
        {
               Random r = new Random(System.currentTimeMillis());
               for(int t=0;t<N;t++)
               {
                  int tp= r.nextInt(N);
                  tempp[t]=Pvalue[tp];
               }  
               for(int j=0;j<20;j++)
              {
                 int c=0;
                  for(int jj=0;jj<N;jj++)
                    if(tempp[jj]>lambda[j])
                       c++;
                  p[j]+=Math.pow(c/((1-lambda[j])*N)-pi0,2);
             } 
        }
        
        for(int i=0;i<20;i++)
        {
            p[i]=p[i]/10000;
        }
        int k=0;
        double min=p[0];
        for(int i=0;i<20;i++)
        { 
            if(min>=p[i])
            {
                min=p[i];
                k=i;
            }
        }       
        pi0=pp[k];
        System.out.println("pi0 = "+pi0);
        return pi0;
    }

       
    public void printf(String file,double pi0)
    {      
        try { 
            BufferedWriter output = new BufferedWriter(new FileWriter(file));
            output.write("pi0= " + pi0 + "\n");
         
            for (int i = 0; i < N; i++) {               
                output.write(proteins.get(i)+"\t"+Y[i] + "\t" + Pvalue[i] + "\t" + fdr[i]+"\n");            
            }
            output.flush();
            output.close();
        } catch (Exception e) { }
    }
    
  
    public static void main(String[] args) 
    {
        double startTime = System.currentTimeMillis();
        Decoy_free_FDR_estimation fdr_lasso = new Decoy_free_FDR_estimation();
     //   fdr_lasso.initial("data\\Human\\Human_peptide_idenfication.txt","data\\Human\\LP.txt");
     // fdr_lasso.initial("data\\DME\\DME_peptide_identification.txt","data\\DME\\Pro.txt");
    //   fdr_lasso.initial("data\\MS\\MSPresso_Human_peptide_idenfication.txt","data\\MS\\LP.txt");
     //   fdr_lasso.initial("data\\18mixture\\181peptide.txt","data\\18mixture\\181pro.txt");
      fdr_lasso.initial("data\\sigma\\sigma1.txt","data\\sigma\\sigma1pro.txt");
      //  fdr_lasso.initial("data\\yeast\\yeastpeptide.txt","data\\yeast\\yeastpro.txt");
        Pvalue = fdr_lasso.compute_pvalue(10000,Y);
        double pi0 = fdr_lasso.choose_pi0(Pvalue);
        fdr = fdr_lasso.estimate_FDR(pi0, Pvalue);
   //     fdr_lasso.printf("result\\yeastpro.xls", pi0);
          fdr_lasso.printf("sigma11.xls", pi0);
       double endTime = System.currentTimeMillis();
       double running_time = (endTime - startTime) / (double) 1000;
       System.out.println("Running Time:" + running_time);
    }
}

class Edge
{
    int first;
    int second;
}

class Protein {
       String name;
       double pro;

       Protein(String n, double p)
       {
                name = n;
                pro = p;
       }
}

